import asyncio
import json
import logging
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack
import time
import fractions
from typing import Tuple, Dict, Optional, Set, Union
from frames_inference import LipModel
import os
import uuid
import wave
from queue import Queue
import threading
from frames_webrtc import HumanPlayer
import concurrent.futures
import aiohttp


# 配置日志
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

pcs = set()
shared_queue = asyncio.Queue()
lipModels: Dict[str, LipModel] = {}  # sessionid: LipModel

# 配置文件上传和下载的目录
UPLOAD_FOLDER = 'uploads'
OUTPUT_FOLDER = 'outputs'
ALLOWED_EXTENSIONS = {'mp4', 'wav', 'mov', 'avi', 'jpg', 'png'}

# 创建必要的目录
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
os.makedirs('temp', exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

async def synch_lips(request):
    # 检查是否提供了视频和音频文件
    data = await request.post()
    video_file = data.get('video')
    audio_file = data.get('audio')

    if video_file is None or audio_file is None:
        return web.json_response({'error': 'Video and audio files are required'}, status=400)

    if video_file.filename == '' or audio_file.filename == '':
        return web.json_response({'error': 'Invalid file'}, status=400)

    if not allowed_file(video_file.filename) or not allowed_file(audio_file.filename):
        return web.json_response({'error': 'Invalid file type'}, status=400)

    # 生成唯一的请求ID，用于区分不同的请求
    request_id = str(uuid.uuid4())

    # 保存上传的文件
    video_path = os.path.join(UPLOAD_FOLDER, f"{request_id}_video.{video_file.filename.rsplit('.', 1)[1].lower()}")
    output_path = os.path.join(OUTPUT_FOLDER, f"{request_id}_output.mp4")
    pcm_path = os.path.join(UPLOAD_FOLDER, f"{request_id}_audio.pcm")
    wav_path = os.path.join(UPLOAD_FOLDER, f"{request_id}_audio.wav")

    # 保存视频文件
    video_content = video_file.file.read()
    video_size = len(video_content)
    logger.info(f"Received video file size: {video_size} bytes")
    with open(video_path, 'wb') as f:
        f.write(video_content)
    logger.info(f"Saved video file size: {os.path.getsize(video_path)} bytes")

    # 保存音频文件
    audio_content = audio_file.file.read()
    audio_size = len(audio_content)
    logger.info(f"Received audio file size: {audio_size} bytes")
    with open(pcm_path, 'wb') as f:
        f.write(audio_content)
    logger.info(f"Saved audio file size: {os.path.getsize(pcm_path)} bytes")

    # 将PCM转换为WAV
    pcm_to_wav(pcm_path, wav_path)

    await shared_queue.put({
    'request_id': request_id,
    'video_path': video_path,
    'audio_path': wav_path,
    'output_path': output_path
    })

    return web.json_response({'message': 'Processing started', 'request_id': request_id}, status=202)

def pcm_to_wav(pcm_path, wav_path, sample_rate=16000, bit_depth=16, channels=1):
    """
    将PCM文件转换为WAV格式，并在转换完成后删除PCM源文件
    :param pcm_path: PCM文件路径
    :param wav_path: 输出WAV文件路径
    :param sample_rate: 采样率，默认16000 Hz
    :param bit_depth: 位深，默认16位
    :param channels: 声道数，默认1
    """
    try:
        # 读取PCM数据
        with open(pcm_path, 'rb') as pcm_file:
            pcm_data = pcm_file.read()

        # 创建WAV文件
        with wave.open(wav_path, 'wb') as wav_file:
            wav_file.setnchannels(channels)
            wav_file.setsampwidth(bit_depth // 8)
            wav_file.setframerate(sample_rate)
            wav_file.writeframes(pcm_data)

        print(f"PCM文件已成功转换为WAV格式，输出路径为: {wav_path}")

        # 删除PCM源文件
        os.remove(pcm_path)
        print(f"PCM源文件已删除: {pcm_path}")

    except Exception as e:
        print(f"发生错误: {e}")

async def offer(request):
    params = await request.json()
    logging.info("Received offer: %s", params["sdp"])
    offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])

    pc = RTCPeerConnection()
    pcs.add(pc)

    try:
        config = await shared_queue.get()
    except asyncio.QueueEmpty:
        return web.Response(status=404, text='No parameters available')

    request_id = config['request_id']
    video_path = config['video_path']
    audio_path = config['audio_path']
    
    lipModels[request_id] = None

    lip_model = LipModel(face_path=video_path, audio_path=audio_path)
    lipModels[request_id] = lip_model
    
    #TODO：轮询遍历lip_model.on_complete这个bool值是否为true，如果为true那么关闭连接

    @pc.on("iceconnectionstatechange")
    async def on_ice_state_change():
        logging.info("ICE connection state is %s", pc.iceConnectionState)
        if pc.iceConnectionState == "failed":
            await pc.close()
            pcs.discard(pc)

    player = HumanPlayer(lip_model, audio_path)
    pc.addTrack(player.video)

    await pc.setRemoteDescription(offer)
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    if pc.localDescription is None:
        return web.Response(status=500, text='Failed to create answer')

    processing_thread = threading.Thread(target=lip_model.process, daemon=True)
    processing_thread.start()

    # 启动轮询任务
    asyncio.create_task(poll_on_complete(pc, lip_model))

    response = {
        "sdp": pc.localDescription.sdp,
        "type": pc.localDescription.type
    }
    return web.json_response(response)


async def poll_on_complete(pc, lip_model):
    while True:
        if lip_model.on_complete:
            await close_webrtc(pc)
            break
        await asyncio.sleep(0.5)  # 每秒检查一次

async def close_webrtc(pc):
    if pc:
        # 清理资源
        pc.onicecandidate = None
        pc.ontrack = None
        await pc.close()
        pcs.discard(pc)
        logger.info("WebRTC 连接已关闭并清理资源")


async def getStatus(request):
    params = await request.json()
    request_id = params.get('request_id')
    if not request_id:
        return web.json_response({'error': 'Request ID is required'}, status=400)

    lip_model = lipModels.get(request_id)
    if not lip_model:
        return web.json_response({'error': 'Request ID not found'}, status=404)

    status = {
        'processing': not lip_model.on_complete,
        'completed': lip_model.on_complete
    }
    return web.json_response(status)

async def on_shutdown(app):
    # 程序退出时关闭所有 peer connection
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()

# 创建 aiohttp 应用
app = web.Application(client_max_size=100 * 1024 * 1024)
app.router.add_post("/synch_lips", synch_lips)
app.router.add_post("/offer", offer)
app.router.add_post("/GetStatus",getStatus)
app.on_shutdown.append(on_shutdown)

# 启动 aiohttp 应用
if __name__ == "__main__":
    web.run_app(app, port=5000)